#!/usr/bin/python3
# coding=utf-8
#
# Copyright (C) 2023-2024. Huawei Technologies Co., Ltd. All rights reserved.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
# ===============================================================================

import numpy as np
import torch
import torch.nn.functional as F
import os

def im2col_torch(x, kernel_size, stride=1, padding=0):
    """
    PyTorch实现的im2col，使用F.unfold
    x: [N, C, H, W] tensor
    return: [N, C*kH*kW, L] where L = out_H * out_W
    """
    if isinstance(kernel_size, int):
        kernel_size = (kernel_size, kernel_size)
    
    # 使用PyTorch的unfold函数
    col = F.unfold(x, kernel_size, stride=stride, padding=padding)
    return col

def gen_golden_data_im2col_torch():
    input_shape = (1, 1, 9, 257)
    input_data_np = np.random.uniform(low=-1.0, high=1.0, size=input_shape).astype(np.float32)
    input_data = torch.from_numpy(input_data_np)
    
    # im2col参数
    kernel_size = (2, 2)  # 卷积核大小
    stride = 1           # 步长
    padding = 0          # 填充
    
    # 计算im2col
    golden_tensor = im2col_torch(input_data, kernel_size, stride, padding)
    golden = golden_tensor.numpy().astype(np.float32)
    
    # 创建目录
    os.makedirs("input", exist_ok=True)
    os.makedirs("output", exist_ok=True)
    
    # 保存输入和输出
    input_data_np.tofile("input/input_x.bin")
    golden.tofile("output/golden.bin")

if __name__ == "__main__":
    # 使用PyTorch实现
    gen_golden_data_im2col_torch()